{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agents\n",
    "\n",
    "AutoGen AgentChat provides a set of preset Agents, each with variations in how an agent might respond to messages.\n",
    "All agents share the following attributes and methods:\n",
    "\n",
    "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.name`: The unique name of the agent.\n",
    "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.description`: The description of the agent in text.\n",
    "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.run`: The method that runs the agent given a task as a string or a list of messages, and returns a {py:class}`~autogen_agentchat.base.TaskResult`. **Agents are expected to be stateful and this method is expected to be called with new messages, not complete history**.\n",
    "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.run_stream`: Same as {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` but returns an iterator of messages that subclass {py:class}`~autogen_agentchat.messages.BaseAgentEvent` or {py:class}`~autogen_agentchat.messages.BaseChatMessage` followed by a {py:class}`~autogen_agentchat.base.TaskResult` as the last item.\n",
    "\n",
    "See {py:mod}`autogen_agentchat.messages` for more information on AgentChat message types."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assistant Agent\n",
    "\n",
    "{py:class}`~autogen_agentchat.agents.AssistantAgent` is a built-in agent that\n",
    "uses a language model and has the ability to use tools.\n",
    "\n",
    "```{warning}\n",
    "{py:class}`~autogen_agentchat.agents.AssistantAgent` is a \"kitchen sink\" agent\n",
    "for prototyping and educational purpose -- it is very general.\n",
    "Make sure you read the documentation and implementation to understand the design choices.\n",
    "Once you fully understand the design, you may want to implement your own agent.\n",
    "See [Custom Agent](../custom-agents.ipynb).\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_agentchat.messages import StructuredMessage\n",
    "from autogen_agentchat.ui import Console\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a tool that searches the web for information.\n",
    "# For simplicity, we will use a mock function here that returns a static string.\n",
    "async def web_search(query: str) -> str:\n",
    "    \"\"\"Find information on the web\"\"\"\n",
    "    return \"AutoGen is a programming framework for building multi-agent applications.\"\n",
    "\n",
    "\n",
    "# Create an agent that uses the OpenAI GPT-4o model.\n",
    "model_client = OpenAIChatCompletionClient(\n",
    "    model=\"gpt-4.1-nano\",\n",
    "    # api_key=\"YOUR_API_KEY\",\n",
    ")\n",
    "agent = AssistantAgent(\n",
    "    name=\"assistant\",\n",
    "    model_client=model_client,\n",
    "    tools=[web_search],\n",
    "    system_message=\"Use tools to solve tasks.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting Result\n",
    "\n",
    "We can use the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` method to get the agent run on a given task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[TextMessage(source='user', models_usage=None, metadata={}, content='Find information on AutoGen', type='TextMessage'), ToolCallRequestEvent(source='assistant', models_usage=RequestUsage(prompt_tokens=61, completion_tokens=16), metadata={}, content=[FunctionCall(id='call_703i17OLXfztkuioUbkESnea', arguments='{\"query\":\"AutoGen\"}', name='web_search')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='assistant', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='AutoGen is a programming framework for building multi-agent applications.', name='web_search', call_id='call_703i17OLXfztkuioUbkESnea', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='assistant', models_usage=None, metadata={}, content='AutoGen is a programming framework for building multi-agent applications.', type='ToolCallSummaryMessage')]\n"
     ]
    }
   ],
   "source": [
    "# Use asyncio.run(agent.run(...)) when running in a script.\n",
    "result = await agent.run(task=\"Find information on AutoGen\")\n",
    "print(result.messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The call to the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` method\n",
    "returns a {py:class}`~autogen_agentchat.base.TaskResult`\n",
    "with the list of messages in the {py:attr}`~autogen_agentchat.base.TaskResult.messages` attribute,\n",
    "which stores the agent's \"thought process\" as well as the final response.\n",
    "\n",
    "```{note}\n",
    "It is important to note that {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`\n",
    "will update the internal state of the agent -- it will add the messages to the agent's\n",
    "message history. You can also call {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`\n",
    "without a task to get the agent to generate responses given its current state.\n",
    "```\n",
    "\n",
    "```{note}\n",
    "Unlike in v0.2 AgentChat, the tools are executed by the same agent directly within\n",
    "the same call to {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`.\n",
    "By default, the agent will return the result of the tool call as the final response.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-Modal Input\n",
    "\n",
    "The {py:class}`~autogen_agentchat.agents.AssistantAgent` can handle multi-modal input\n",
    "by providing the input as a {py:class}`~autogen_agentchat.messages.MultiModalMessage`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"\"/>"
      ],
      "text/plain": [
       "<autogen_core._image.Image at 0x7f7089c83710>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from io import BytesIO\n",
    "\n",
    "import PIL\n",
    "import requests\n",
    "from autogen_agentchat.messages import MultiModalMessage\n",
    "from autogen_core import Image\n",
    "\n",
    "# Create a multi-modal message with random image and text.\n",
    "pil_image = PIL.Image.open(BytesIO(requests.get(\"https://picsum.photos/300/200\").content))\n",
    "img = Image(pil_image)\n",
    "multi_modal_message = MultiModalMessage(content=[\"Can you describe the content of this image?\", img], source=\"user\")\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The image depicts a scenic mountain landscape under a clear blue sky. There are several rugged mountain peaks in the background, with some clouds scattered across the sky. In the valley below, there is a body of water, possibly a lake or river, surrounded by greenery. The overall scene conveys a sense of natural beauty and tranquility.\n"
     ]
    }
   ],
   "source": [
    "# Use asyncio.run(...) when running in a script.\n",
    "result = await agent.run(task=multi_modal_message)\n",
    "print(result.messages[-1].content)  # type: ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming Messages\n",
    "\n",
    "We can also stream each message as it is generated by the agent by using the\n",
    "{py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream` method,\n",
    "and use {py:class}`~autogen_agentchat.ui.Console` to print the messages\n",
    "as they appear to the console."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- TextMessage (user) ----------\n",
      "Find information on AutoGen\n",
      "---------- ToolCallRequestEvent (assistant) ----------\n",
      "[FunctionCall(id='call_HOTRhOzXCBm0zSqZCFbHD7YP', arguments='{\"query\":\"AutoGen\"}', name='web_search')]\n",
      "[Prompt tokens: 61, Completion tokens: 16]\n",
      "---------- ToolCallExecutionEvent (assistant) ----------\n",
      "[FunctionExecutionResult(content='AutoGen is a programming framework for building multi-agent applications.', name='web_search', call_id='call_HOTRhOzXCBm0zSqZCFbHD7YP', is_error=False)]\n",
      "---------- ToolCallSummaryMessage (assistant) ----------\n",
      "AutoGen is a programming framework for building multi-agent applications.\n",
      "---------- Summary ----------\n",
      "Number of messages: 4\n",
      "Finish reason: None\n",
      "Total prompt tokens: 61\n",
      "Total completion tokens: 16\n",
      "Duration: 0.52 seconds\n"
     ]
    }
   ],
   "source": [
    "async def assistant_run_stream() -> None:\n",
    "    # Option 1: read each message from the stream (as shown in the previous example).\n",
    "    # async for message in agent.run_stream(task=\"Find information on AutoGen\"):\n",
    "    #     print(message)\n",
    "\n",
    "    # Option 2: use Console to print all messages as they appear.\n",
    "    await Console(\n",
    "        agent.run_stream(task=\"Find information on AutoGen\"),\n",
    "        output_stats=True,  # Enable stats printing.\n",
    "    )\n",
    "\n",
    "\n",
    "# Use asyncio.run(assistant_run_stream()) when running in a script.\n",
    "await assistant_run_stream()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream` method\n",
    "returns an asynchronous generator that yields each message generated by the agent,\n",
    "followed by a {py:class}`~autogen_agentchat.base.TaskResult` as the last item.\n",
    "\n",
    "From the messages, you can observe that the assistant agent utilized the `web_search` tool to\n",
    "gather information and responded based on the search results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Tools and Workbench\n",
    "\n",
    "Large Language Models (LLMs) are typically limited to generating text or code responses. \n",
    "However, many complex tasks benefit from the ability to use external tools that perform specific actions,\n",
    "such as fetching data from APIs or databases.\n",
    "\n",
    "To address this limitation, modern LLMs can now accept a list of available tool schemas \n",
    "(descriptions of tools and their arguments) and generate a tool call message. \n",
    "This capability is known as **Tool Calling** or **Function Calling** and \n",
    "is becoming a popular pattern in building intelligent agent-based applications.\n",
    "Refer to the documentation from [OpenAI](https://platform.openai.com/docs/guides/function-calling) \n",
    "and [Anthropic](https://docs.anthropic.com/en/docs/build-with-claude/tool-use) for more information about tool calling in LLMs.\n",
    "\n",
    "In AgentChat, the {py:class}`~autogen_agentchat.agents.AssistantAgent` can use tools to perform specific actions.\n",
    "The `web_search` tool is one such tool that allows the assistant agent to search the web for information.\n",
    "A single custom tool can be a Python function or a subclass of the {py:class}`~autogen_core.tools.BaseTool`.\n",
    "\n",
    "On the other hand, a {py:class}`~autogen_core.tools.Workbench` is a collection of tools that share state and resources.\n",
    "\n",
    "```{note}\n",
    "For how to use model clients directly with tools and workbench, refer to the [Tools](../../core-user-guide/components/tools.ipynb)\n",
    "and [Workbench](../../core-user-guide/components/workbench.ipynb) sections\n",
    "in the Core User Guide.\n",
    "```\n",
    "\n",
    "By default, when {py:class}`~autogen_agentchat.agents.AssistantAgent` executes a tool,\n",
    "it will return the tool's output as a string in {py:class}`~autogen_agentchat.messages.ToolCallSummaryMessage` in its response.\n",
    "If your tool does not return a well-formed string in natural language, you\n",
    "can add a reflection step to have the model summarize the tool's output,\n",
    "by setting the `reflect_on_tool_use=True` parameter in the {py:class}`~autogen_agentchat.agents.AssistantAgent` constructor.\n",
    "\n",
    "### Built-in Tools and Workbench\n",
    "\n",
    "AutoGen Extension provides a set of built-in tools that can be used with the Assistant Agent.\n",
    "Head over to the [API documentation](../../../reference/index.md) for all the available tools\n",
    "under the `autogen_ext.tools` namespace. For example, you can find the following tools:\n",
    "\n",
    "- {py:mod}`~autogen_ext.tools.graphrag`: Tools for using GraphRAG index.\n",
    "- {py:mod}`~autogen_ext.tools.http`: Tools for making HTTP requests.\n",
    "- {py:mod}`~autogen_ext.tools.langchain`: Adaptor for using LangChain tools.\n",
    "- {py:mod}`~autogen_ext.tools.mcp`: Tools and workbench for using Model Chat Protocol (MCP) servers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function Tool\n",
    "\n",
    "The {py:class}`~autogen_agentchat.agents.AssistantAgent` automatically\n",
    "converts a Python function into a {py:class}`~autogen_core.tools.FunctionTool`\n",
    "which can be used as a tool by the agent and automatically generates the tool schema\n",
    "from the function signature and docstring.\n",
    "\n",
    "The `web_search_func` tool is an example of a function tool.\n",
    "The schema is automatically generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'web_search_func',\n",
       " 'description': 'Find information on the web',\n",
       " 'parameters': {'type': 'object',\n",
       "  'properties': {'query': {'description': 'query',\n",
       "    'title': 'Query',\n",
       "    'type': 'string'}},\n",
       "  'required': ['query'],\n",
       "  'additionalProperties': False},\n",
       " 'strict': False}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from autogen_core.tools import FunctionTool\n",
    "\n",
    "\n",
    "# Define a tool using a Python function.\n",
    "async def web_search_func(query: str) -> str:\n",
    "    \"\"\"Find information on the web\"\"\"\n",
    "    return \"AutoGen is a programming framework for building multi-agent applications.\"\n",
    "\n",
    "\n",
    "# This step is automatically performed inside the AssistantAgent if the tool is a Python function.\n",
    "web_search_function_tool = FunctionTool(web_search_func, description=\"Find information on the web\")\n",
    "# The schema is provided to the model during AssistantAgent's on_messages call.\n",
    "web_search_function_tool.schema"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Context Protocol (MCP) Workbench\n",
    "\n",
    "The {py:class}`~autogen_agentchat.agents.AssistantAgent` can also use tools that are\n",
    "served from a Model Context Protocol (MCP) server\n",
    "using {py:func}`~autogen_ext.tools.mcp.McpWorkbench`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Seattle is a major city located in the state of Washington, United States. It was founded on November 13, 1851, and incorporated as a town on January 14, 1865, and later as a city on December 2, 1869. The city is named after Chief Seattle. It covers an area of approximately 142 square miles, with a population of around 737,000 as of the 2020 Census, and an estimated 755,078 residents in 2023. Seattle is known by nicknames such as The Emerald City, Jet City, and Rain City, and has mottos including The City of Flowers and The City of Goodwill. The city operates under a mayor–council government system, with Bruce Harrell serving as mayor. Key landmarks include the Space Needle, Pike Place Market, Amazon Spheres, and the Seattle Great Wheel. It is situated on the U.S. West Coast, with a diverse urban and metropolitan area that extends to a population of over 4 million in the greater metropolitan region.\n"
     ]
    }
   ],
   "source": [
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_agentchat.messages import TextMessage\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams\n",
    "\n",
    "# Get the fetch tool from mcp-server-fetch.\n",
    "fetch_mcp_server = StdioServerParams(command=\"uvx\", args=[\"mcp-server-fetch\"])\n",
    "\n",
    "# Create an MCP workbench which provides a session to the mcp server.\n",
    "async with McpWorkbench(fetch_mcp_server) as workbench:  # type: ignore\n",
    "    # Create an agent that can use the fetch tool.\n",
    "    model_client = OpenAIChatCompletionClient(model=\"gpt-4.1-nano\")\n",
    "    fetch_agent = AssistantAgent(\n",
    "        name=\"fetcher\", model_client=model_client, workbench=workbench, reflect_on_tool_use=True\n",
    "    )\n",
    "\n",
    "    # Let the agent fetch the content of a URL and summarize it.\n",
    "    result = await fetch_agent.run(task=\"Summarize the content of https://en.wikipedia.org/wiki/Seattle\")\n",
    "    assert isinstance(result.messages[-1], TextMessage)\n",
    "    print(result.messages[-1].content)\n",
    "\n",
    "    # Close the connection to the model client.\n",
    "    await model_client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agent as a Tool\n",
    "\n",
    "Any {py:class}`~autogen_agentchat.agents.BaseChatAgent` can be used as a tool\n",
    "by wrapping it in a {py:class}`~autogen_agentchat.tools.AgentTool`.\n",
    "This allows for a dynamic, model-driven multi-agent workflow where\n",
    "the agent can call other agents as tools to solve tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel Tool Calls\n",
    "\n",
    "Some models support parallel tool calls, which can be useful for tasks that require multiple tools to be called simultaneously.\n",
    "By default, if the model client produces multiple tool calls, {py:class}`~autogen_agentchat.agents.AssistantAgent`\n",
    "will call the tools in parallel.\n",
    "\n",
    "You may want to disable parallel tool calls when the tools have side effects that may interfere with each other, or,\n",
    "when agent behavior needs to be consistent across different models.\n",
    "This should be done at the model client level.\n",
    "\n",
    "```{important}\n",
    "When using {py:class}`~autogen_agentchat.tools.AgentTool` or {py:class}`~autogen_agentchat.tools.TeamTool`,\n",
    "you **must** disable parallel tool calls to avoid concurrency issues.\n",
    "These tools cannot run concurrently as agents and teams maintain internal state\n",
    "that would conflict with parallel execution.\n",
    "```\n",
    "\n",
    "For {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient` and {py:class}`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`,\n",
    "set `parallel_tool_calls=False` to disable parallel tool calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_client_no_parallel_tool_call = OpenAIChatCompletionClient(\n",
    "    model=\"gpt-4o\",\n",
    "    parallel_tool_calls=False,  # type: ignore\n",
    ")\n",
    "agent_no_parallel_tool_call = AssistantAgent(\n",
    "    name=\"assistant\",\n",
    "    model_client=model_client_no_parallel_tool_call,\n",
    "    tools=[web_search],\n",
    "    system_message=\"Use tools to solve tasks.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tool Iterations\n",
    "\n",
    "One model call followed by one tool call or parallel tool calls\n",
    "is a single tool iteration.\n",
    "By default, the {py:class}`~autogen_agentchat.agents.AssistantAgent` will\n",
    "execute at most one iteration.\n",
    "\n",
    "The agent can be configured to execute multiple iterations until the model\n",
    "stops generating tool calls or the maximum number of iterations is reached.\n",
    "You can control the maximum number of iterations by setting the `max_tool_iterations` parameter\n",
    "in the {py:class}`~autogen_agentchat.agents.AssistantAgent` constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_loop = AssistantAgent(\n",
    "    name=\"assistant_loop\",\n",
    "    model_client=model_client_no_parallel_tool_call,\n",
    "    tools=[web_search],\n",
    "    system_message=\"Use tools to solve tasks.\",\n",
    "    max_tool_iterations=10,  # At most 10 iterations of tool calls before stopping the loop.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured Output\n",
    "\n",
    "Structured output allows models to return structured JSON text with pre-defined schema\n",
    "provided by the application. Different from JSON-mode, the schema can be provided\n",
    "as a [Pydantic BaseModel](https://docs.pydantic.dev/latest/concepts/models/)\n",
    "class, which can also be used to validate the output.\n",
    "\n",
    "Once you specify the base model class in the `output_content_type` parameter\n",
    "of the {py:class}`~autogen_agentchat.agents.AssistantAgent` constructor,\n",
    "the agent will respond with a {py:class}`~autogen_agentchat.messages.StructuredMessage`\n",
    "whose `content`'s type is the type of the base model class.\n",
    "\n",
    "This way, you can integrate agent's response directly into your application\n",
    "and use the model's output as a structured object.\n",
    "\n",
    "```{note}\n",
    "When the `output_content_type` is set, it by default requires the agent to reflect on the tool use\n",
    "and return the a structured output message based on the tool call result.\n",
    "You can disable this behavior by setting `reflect_on_tool_use=False` explictly.\n",
    "```\n",
    "\n",
    "Structured output is also useful for incorporating Chain-of-Thought\n",
    "reasoning in the agent's responses.\n",
    "See the example below for how to use structured output with the assistant agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- user ----------\n",
      "I am happy.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- assistant ----------\n",
      "{\n",
      "  \"thoughts\": \"The user explicitly states they are happy.\",\n",
      "  \"response\": \"happy\"\n",
      "}\n",
      "Thought:  The user explicitly states they are happy.\n",
      "Response:  happy\n"
     ]
    }
   ],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "# The response format for the agent as a Pydantic base model.\n",
    "class AgentResponse(BaseModel):\n",
    "    thoughts: str\n",
    "    response: Literal[\"happy\", \"sad\", \"neutral\"]\n",
    "\n",
    "\n",
    "# Create an agent that uses the OpenAI GPT-4o model.\n",
    "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n",
    "agent = AssistantAgent(\n",
    "    \"assistant\",\n",
    "    model_client=model_client,\n",
    "    system_message=\"Categorize the input as happy, sad, or neutral following the JSON format.\",\n",
    "    # Define the output content type of the agent.\n",
    "    output_content_type=AgentResponse,\n",
    ")\n",
    "\n",
    "result = await Console(agent.run_stream(task=\"I am happy.\"))\n",
    "\n",
    "# Check the last message in the result, validate its type, and print the thoughts and response.\n",
    "assert isinstance(result.messages[-1], StructuredMessage)\n",
    "assert isinstance(result.messages[-1].content, AgentResponse)\n",
    "print(\"Thought: \", result.messages[-1].content.thoughts)\n",
    "print(\"Response: \", result.messages[-1].content.response)\n",
    "await model_client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming Tokens\n",
    "\n",
    "You can stream the tokens generated by the model client by setting `model_client_stream=True`.\n",
    "This will cause the agent to yield {py:class}`~autogen_agentchat.messages.ModelClientStreamingChunkEvent` messages\n",
    "in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream`.\n",
    "\n",
    "The underlying model API must support streaming tokens for this to work.\n",
    "Please check with your model provider to see if this is supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "source='user' models_usage=None metadata={} content='Name two cities in South America' type='TextMessage'\n",
      "source='assistant' models_usage=None metadata={} content='Two' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' cities' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' in' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' South' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' America' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' are' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' Buenos' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' Aires' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' in' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' Argentina' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' and' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' São' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' Paulo' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' in' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content=' Brazil' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=None metadata={} content='.' type='ModelClientStreamingChunkEvent'\n",
      "source='assistant' models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0) metadata={} content='Two cities in South America are Buenos Aires in Argentina and São Paulo in Brazil.' type='TextMessage'\n",
      "messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Name two cities in South America', type='TextMessage'), TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=0, completion_tokens=0), metadata={}, content='Two cities in South America are Buenos Aires in Argentina and São Paulo in Brazil.', type='TextMessage')] stop_reason=None\n"
     ]
    }
   ],
   "source": [
    "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n",
    "\n",
    "streaming_assistant = AssistantAgent(\n",
    "    name=\"assistant\",\n",
    "    model_client=model_client,\n",
    "    system_message=\"You are a helpful assistant.\",\n",
    "    model_client_stream=True,  # Enable streaming tokens.\n",
    ")\n",
    "\n",
    "# Use an async function and asyncio.run() in a script.\n",
    "async for message in streaming_assistant.run_stream(task=\"Name two cities in South America\"):  # type: ignore\n",
    "    print(message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see the streaming chunks in the output above.\n",
    "The chunks are generated by the model client and are yielded by the agent as they are received.\n",
    "The final response, the concatenation of all the chunks, is yielded right after the last chunk."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Model Context\n",
    "\n",
    "{py:class}`~autogen_agentchat.agents.AssistantAgent` has a `model_context`\n",
    "parameter that can be used to pass in a {py:class}`~autogen_core.model_context.ChatCompletionContext`\n",
    "object. This allows the agent to use different model contexts, such as\n",
    "{py:class}`~autogen_core.model_context.BufferedChatCompletionContext` to\n",
    "limit the context sent to the model.\n",
    "\n",
    "By default, {py:class}`~autogen_agentchat.agents.AssistantAgent` uses\n",
    "the {py:class}`~autogen_core.model_context.UnboundedChatCompletionContext`\n",
    "which sends the full conversation history to the model. To limit the context\n",
    "to the last `n` messages, you can use the {py:class}`~autogen_core.model_context.BufferedChatCompletionContext`.\n",
    "To limit the context by token count, you can use the\n",
    "{py:class}`~autogen_core.model_context.TokenLimitedChatCompletionContext`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_core.model_context import BufferedChatCompletionContext\n",
    "\n",
    "# Create an agent that uses only the last 5 messages in the context to generate responses.\n",
    "agent = AssistantAgent(\n",
    "    name=\"assistant\",\n",
    "    model_client=model_client,\n",
    "    tools=[web_search],\n",
    "    system_message=\"Use tools to solve tasks.\",\n",
    "    model_context=BufferedChatCompletionContext(buffer_size=5),  # Only use the last 5 messages in the context.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other Preset Agents\n",
    "\n",
    "The following preset agents are available:\n",
    "\n",
    "- {py:class}`~autogen_agentchat.agents.UserProxyAgent`: An agent that takes user input returns it as responses.\n",
    "- {py:class}`~autogen_agentchat.agents.CodeExecutorAgent`: An agent that can execute code.\n",
    "- {py:class}`~autogen_ext.agents.openai.OpenAIAssistantAgent`: An agent that is backed by an OpenAI Assistant, with ability to use custom tools.\n",
    "- {py:class}`~autogen_ext.agents.web_surfer.MultimodalWebSurfer`: A multi-modal agent that can search the web and visit web pages for information.\n",
    "- {py:class}`~autogen_ext.agents.file_surfer.FileSurfer`: An agent that can search and browse local files for information.\n",
    "- {py:class}`~autogen_ext.agents.video_surfer.VideoSurfer`: An agent that can watch videos for information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Step\n",
    "\n",
    "Having explored the usage of the {py:class}`~autogen_agentchat.agents.AssistantAgent`, we can now proceed to the next section to learn about the teams feature in AgentChat.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- ## CodingAssistantAgent\n",
    "\n",
    "Generates responses (text and code) using an LLM upon receipt of a message. It takes a `system_message` argument that defines or sets the tone for how the agent's LLM should respond. \n",
    "\n",
    "```python\n",
    "\n",
    "writing_assistant_agent = CodingAssistantAgent(\n",
    "    name=\"writing_assistant_agent\",\n",
    "    system_message=\"You are a helpful assistant that solve tasks by generating text responses and code.\",\n",
    "    model_client=model_client,\n",
    ")\n",
    "`\n",
    "\n",
    "We can explore or test the behavior of the agent by sending a message to it using the  {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`  method. \n",
    "\n",
    "```python\n",
    "result = await writing_assistant_agent.on_messages(\n",
    "    messages=[\n",
    "        TextMessage(content=\"What is the weather right now in France?\", source=\"user\"),\n",
    "    ],\n",
    "    cancellation_token=CancellationToken(),\n",
    ")\n",
    "print(result) -->"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
